package com.feriki.msync.pool;

import com.feriki.msync.config.ComplexConnectConfig;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

/**milvus连接池
 * @author luozhida
 * @create 2024-04-30
 */
@Slf4j
public class MilvusConnectionPool {

    private final BlockingQueue<MilvusServiceClient> pool = new LinkedBlockingQueue<>(); ;
    private String poolName;
    public MilvusConnectionPool(String poolName, ComplexConnectConfig milvusConfig, int poolSize){
        this.poolName = poolName;
        init(milvusConfig,poolSize);
    }

    private void init(ComplexConnectConfig milvusConfig, int poolSize){
        int maxPoolSize = 1;
        if(poolSize>0){
            maxPoolSize = poolSize;
        }
        for (int i = 0; i < maxPoolSize; i++) {
            try {
                ConnectParam connectParam = ConnectParam.newBuilder()
                        .withHost(milvusConfig.getHost())
                        .withPort(milvusConfig.getPort())
                        .withAuthorization(milvusConfig.getUser(), milvusConfig.getPassword())
                        .build();
                MilvusServiceClient client = new MilvusServiceClient(connectParam);
                client.withRetry(3);
                pool.offer(client);
            } catch (Exception e) {
                log.error("init milvus connection pool throw exception!",e);
                return;
            }
        }
        log.info(String.format("%s-milvus连接池【%s:%s】初始化完成，共有%s个连接",this.poolName,milvusConfig.getHost(),milvusConfig.getPort(),pool.size()));
    }

    public MilvusServiceClient getConnection() throws InterruptedException {
        return pool.take();
    }

    public void releaseConnection(MilvusServiceClient client) {
        if (client == null) return;
        pool.offer(client);
    }

    public void closeAll() {
        for (MilvusServiceClient client : pool) {
            client.close();
        }
        pool.clear();
        log.info(String.format("%s-milvus连接池释放所有连接完成!",this.poolName));
    }

}
